import threading
import time
from typing import Dict, List, Optional
from datetime import datetime, timedelta
import pandas as pd
from xtquant import xtdata
import logging

class DataPreloader:
    """
    数据预加载器
    在启动时预加载所有需要的历史数据，避免在tick回调中实时下载
    """
    
    def __init__(self, stock_codes: List[str], xtdata):
        self.stock_codes = stock_codes
        self.xtdata = xtdata
        
        # 预加载的数据缓存
        self.preloaded_data = {}
        self.preload_status = {}
        
        # 预加载配置
        self.preload_days = 3  # 预加载最近3天的数据
        self.preload_hours = 2  # 预加载最近2小时的数据
        
        # 交易时间段
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()
        
    def preload_all_data(self):
        """预加载所有股票的数据"""
        logging.info("开始预加载历史数据...")
        
        # 启动预加载线程
        preload_thread = threading.Thread(target=self._preload_worker, daemon=True)
        preload_thread.start()
        
        return preload_thread
        
    def _preload_worker(self):
        """预加载工作线程"""
        try:
            # 预加载tick数据
            self._preload_tick_data()
            
            # 预加载分钟数据
            self._preload_minute_data()
            
            logging.info("数据预加载完成")
            
        except Exception as e:
            logging.error(f"数据预加载异常: {e}")
            
    def _preload_tick_data(self):
        """预加载tick数据"""
        today = datetime.now().strftime("%Y%m%d")
        
        for stock_code in self.stock_codes:
            try:
                logging.info(f"预加载tick数据: {stock_code}")
                
                # 下载当日tick数据
                self.xtdata.download_history_data(
                    stock_code=stock_code,
                    period="tick",
                    start_time=today,
                    end_time=today
                )
                time.sleep(0.5)  # 避免请求过于频繁
                
                # 获取本地数据
                data = self.xtdata.get_local_data(
                    stock_list=[stock_code],
                    period='tick',
                    start_time=today,
                    end_time=today
                )
                
                if data and stock_code in data:
                    self.preloaded_data[f"{stock_code}_tick"] = data[stock_code]
                    self.preload_status[f"{stock_code}_tick"] = "success"
                    logging.info(f"tick数据预加载成功: {stock_code}")
                else:
                    self.preload_status[f"{stock_code}_tick"] = "failed"
                    logging.warning(f"tick数据预加载失败: {stock_code}")
                    
            except Exception as e:
                self.preload_status[f"{stock_code}_tick"] = "error"
                logging.error(f"tick数据预加载异常 {stock_code}: {e}")
                
    def _preload_minute_data(self):
        """预加载分钟数据"""
        today = datetime.now().strftime("%Y%m%d")
        
        for stock_code in self.stock_codes:
            try:
                logging.info(f"预加载分钟数据: {stock_code}")
                
                # 下载当日分钟数据
                self.xtdata.download_history_data(
                    stock_code=stock_code,
                    period="1m",
                    start_time=today,
                    end_time=today
                )
                time.sleep(0.3)
                
                # 获取本地数据
                data = self.xtdata.get_local_data(
                    stock_list=[stock_code],
                    period='1m',
                    start_time=today,
                    end_time=today
                )
                
                if data and stock_code in data:
                    self.preloaded_data[f"{stock_code}_minute"] = data[stock_code]
                    self.preload_status[f"{stock_code}_minute"] = "success"
                    logging.info(f"分钟数据预加载成功: {stock_code}")
                else:
                    self.preload_status[f"{stock_code}_minute"] = "failed"
                    logging.warning(f"分钟数据预加载失败: {stock_code}")
                    
            except Exception as e:
                self.preload_status[f"{stock_code}_minute"] = "error"
                logging.error(f"分钟数据预加载异常 {stock_code}: {e}")
                
    def get_preloaded_tick_data(self, stock_code: str) -> Optional[pd.DataFrame]:
        """获取预加载的tick数据"""
        return self.preloaded_data.get(f"{stock_code}_tick")
        
    def get_preloaded_minute_data(self, stock_code: str) -> Optional[pd.DataFrame]:
        """获取预加载的分钟数据"""
        return self.preloaded_data.get(f"{stock_code}_minute")
        
    def is_data_preloaded(self, stock_code: str, data_type: str = "tick") -> bool:
        """检查数据是否已预加载"""
        key = f"{stock_code}_{data_type}"
        return self.preload_status.get(key) == "success"
        
    def get_preload_status(self) -> Dict[str, str]:
        """获取预加载状态"""
        return self.preload_status.copy()
        
    def refresh_data(self, stock_code: str = None):
        """刷新预加载数据"""
        if stock_code:
            # 刷新指定股票的数据
            if f"{stock_code}_tick" in self.preloaded_data:
                del self.preloaded_data[f"{stock_code}_tick"]
            if f"{stock_code}_minute" in self.preloaded_data:
                del self.preloaded_data[f"{stock_code}_minute"]
                
            # 重新预加载
            self._preload_tick_data()
            self._preload_minute_data()
        else:
            # 刷新所有数据
            self.preloaded_data.clear()
            self.preload_status.clear()
            self._preload_tick_data()
            self._preload_minute_data()
            
    def get_data_stats(self) -> Dict[str, any]:
        """获取数据统计信息"""
        total_stocks = len(self.stock_codes)
        tick_success = sum(1 for status in self.preload_status.values() 
                          if status == "success" and "_tick" in status)
        minute_success = sum(1 for status in self.preload_status.values() 
                            if status == "success" and "_minute" in status)
        
        return {
            "total_stocks": total_stocks,
            "tick_data_loaded": tick_success,
            "minute_data_loaded": minute_success,
            "tick_success_rate": f"{tick_success/total_stocks*100:.1f}%" if total_stocks > 0 else "0%",
            "minute_success_rate": f"{minute_success/total_stocks*100:.1f}%" if total_stocks > 0 else "0%"
        }
